import openai
from logs.logger import log_record
from config import AI_CONFIG


class ChatGPTBot:

    def __init__(self):
        config = AI_CONFIG['chatgpt']
        openai.api_key = config['api_key']
        openai.api_base = config['api_base']
        self.max_tokens = config['max_tokens']
        self.temperature = config['temperature']
        self.default_model = config['default_model']

    def generate(self, prompt, model=None, max_tokens=None, temperature=None):
        model = model or self.default_model
        max_tokens = max_tokens or self.max_tokens
        temperature = temperature or self.temperature

        log_record(prompt, "prompt")

        if model == "text-davinci-003":
            response = openai.Completion.create(
                engine=model,
                prompt=prompt,
                max_tokens=max_tokens,
            )
            log_record(response, "response")
            return response.choices[0].text
        else:
            response = openai.ChatCompletion.create(
                model=model,
                messages=prompt,
                temperature=temperature,
            )
            log_record(response, "response")

            # return response.choices[0].message.get("content", '')
            return response
